# -*- encoding: UTF-8 -*-
# message based kademlia rpc module.
# author: orzp2p core group 2009-6-15
# license: opensource, gplv3 
# warnning: this code has not do any test yet

import math
import random
import time

import orzp2p.core as core
import orzp2p.component.socket as sock

msgbus = core.msgbus()
sockid = sock.id

class data(object):
    '''nothing, empty class object'''
    pass
    
class peer(object):
    '''peer:(160bit, ip, port)'''
    def __init__(self, peerid, ip, port):
        self.peerid = peerid
        self.ip = ip
        self.port = port
    
class dhterror(Exception):
    '''2:find a pair not found 1: cache full, max sessions reached'''
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)

class dht(object):
    def __init__(self, localpeer = None, peers = [[]] * 160, dhtcache={}, \
            overcache={}, bucketsecond=3600, k=20, a=3):
        '''args:localpeer, peers, dhtcache, bucketsecond=3600, k=20'''
        random.seed(peers)
        if localpeer is None:
            localpeer = peer(random.randint(0, 2**160 - 1), '127.0.0.1', 10101)
        self.localpeer = localpeer
        self.peers = peers
        self.dhtcache = dhtcache
        self.overcache = overcache
        self.bucketsecond = bucketsecond
        # peers:[[peer, ...], [peer, ...], ...]
        self.k = k
        self.a = a
        self.buckettimeout = [0] * 160
        self.cmdlist = {'peercome':self.peercome, 
                            'findnode':self.findnode, 
                            'findvalue':self.findvalue,
                            'store':self.store,
                            'nodelookup':self.nodelookup,
                            'respondlook':self.respondlook,
                            'findapair':self.findapair,
                            'respondvalue':self.respondvalue,
                            'reflashkb':self.reflashkb,
                            'pong':self.pong,
                            'ping':self.ping,
                            'storapair':self.storapair,
                            'joinus':self.joinus,
                            'overcaching':self.overcaching}
        self.sessions = {}
        self.requested = {}
        self.result= {}
        
    def invoke(self, data):
        '''1 arg, data = data'''
        self.cmdlist[data.cmd](data)
        
    def peercome(self, data):
        '''1 arg, data.comid .peer'''
        bklength = (1, 2, 4, 8, 16)
        bucket_no = math.floor(math.log(self.localpeer.peerid ^ \
                                                data.peer.peerid, 2))
        if data.peer in peers[bucket_no]:
            self.peers[bucket_no].remove(data.peer)
            self.peers[bucket_no].append(data.peer)
        if bucket_no in (0, 1, 2, 3, 4):
            bucket_maxlength = bklength[bucket_no]
        else: 
            bucket_maxlength = self.k
        data.peer =self.peers[bucket_no][0]
        if len(self.peers[bucket_no]) < bucket_maxlength:
            self.peers[bucket_no].append(data.peer)
        elif not self.ping(data):
            del self.peers[bucket_no][0]
            self.peers[bucket_no].append(data.peer)
        else:
            self.peers[bucket_no].append(self.peers[bucket_no].pop(0))
        self.buckettimeout[bucket_no] = time.time() + self.bucketsecond
    
    def findnode(self, data):
        '''1 arg, data.findid .comid .cmd .from_peerid
        .a .session'''
        if len(self.peers) > 0:
            r0 = reduce(lambda x,y:x+y, self.peers)
        else:
            r0=[]
        if len(self.sessions) > 0:
            r1 = reduce(lambda x,y:x+y, self.sessions)
        else:
            r1=[]
        if len(self.requested) > 0:
            r2 = reduce(lambda x,y:x+y, self.requested)
        else:
            r2=[]
        data.result = [set(r0) | set(r1) | set(r2)]
        data.result.sort(cmp = lambda x,y : cmp(x.peerid ^ data.findid, \
                                            y.peerid ^ data.findid))
        data.result = data.result[:self.k]
        if data.from_peerid != self.localpeer.peerid:
            data.cmd = 'respondlook'
            msgbus.put((sockid, data), 'first')
        else:
            return data.result

    def findvalue(self, data):
        '''1 arg, data(value's key, from_peerid, session, a)'''
        peerid = self.localpeer.peerid
        if self.dhtcache.get[data.key] is not None:
            data.result = (1, peerid, self.dhtcache.get[data.key])
        else:
            data.result = (0, peerid, self.findnode(data))
        if data.from_peerid != peerid:
            data.cmd = 'respondvalue'
            msgbus.put((sockid, data), 'first')
        else:
            return data.result

    def store(self, data):
        '''1 arg, data.cmd kv dict(key=value)'''
        # find the node nearest to key
        bucket_no = math.floor(math.log(self.localpeer.peerid ^ \
                                                data.kv[key] , 2))
        data.findid = data.kv[key]
        data.from_peerid = self.localpeer.peerid
        temp = self.findnode(data)[0]
        # the nearst is temp node,
        # count node number between key-node and localpeer
        i = 0
        for p in self.peers[bucket_no]:
            i +=1
            if temp == p:
                break
        n = 0
        for j in xrange(bucket_no):
            if j < 5:
                n += 2 ** j  - 1
            else:
                n += self.k
        n += i - 1
        self.dhtcache.update(data.kv)
        keeptime = 1/(2 ** n) * self.bucketsecond
        self.overcache[data.kv[key]] = time.time() + keeptime

    def nodelookup(self, data):
        '''1 arg, data.findid .comid .cmd'''
        if len(self.sessions) >= 320:
#			try:
            raise dhterror(1)
#			except dhterror, e:
#				return e.value
        else:
            data.from_peerid = self.localpeer.peerid
            random.seed(self.peers)
            data.session = random.randint(0, 2**160 - 1)
            data.cmd = 'findnode'
            near = self.findnode(data)
            self.sessions[data.session] = near
            temp = near[:self.a]
        for p in temp:
            msgbus.put((sockid, data), 'first')
        self.requested[data.session] = temp
        # from self.result[data.session] get recv data
            
    def respondlook(self, data):
        '''1 arg, data.comid .findid .from_peerid .result .session)'''
        step_result = data.result
        session = data.session
        data.result = []
        near = [set(self.findnode(data)) | set(step_result)]
        near.sort(cmp = lambda x,y : cmp(x.peerid ^ data.findid, \
                                                y.peerid ^ data.findid))
        near=near[:self.k]
        if self.sessions[session] == near and \
            self.requested[session] == near:
            del self.requested[session]
            del self.sessions[session]
            self.result[session] = near
            return
        else:
            self.sessions[near] = near
            unrequest = set(near) - set(self.requested[session])
            unrequest.sort(cmp = lambda x, y : cmp(x.peerid ^ \
                                data.findid, y.peerid ^ data.findid))
            temp = unrequest[:self.a]
        for p in temp:
            msgbus.put((sockid, data), 'first')
        self.requested[session] += temp
        self.requested[session].sort(cmp = lambda x, y : cmp(x.peerid ^ \
                                data.findid, ypeerid ^ data.findid))
        self.requested[session] = self.requested[session][:self.k]

    def findapair(self, data):
        '''1 arg, data. value's key .comid'''
        data.from_peerid = self.localpeer.peerid
        r = self.findvalue(data)
        if r[0]:
            return r[1]
        if len(self.sessions)  >= 320:
            #try:
            raise dhterror(1)
            #except dhterror, e:
            #	return e.value
        else:
            random.seed(self.peers)
            data.session = random.randint(0, 2**160 - 1)
            data.cmd = 'findvalue'
            self.sessions[data.session] = r[1]
            temp = r[1][:self.a]
        for p in temp:
            msgbus.put((sockid, data), 'first')
        self.requested[data.session] = temp
        # from self.result[data.session] get recv data
        
    def respondvalue(self, data):
        '''1 arg, data(comid, key, from_peerid result, session)'''
        step_result = data.result
        session = data.session
        for r in step_result:
            if r[0]:
                self.result[session]=r[2]
                hasvalue = r[1]
                if hasvalue in self.requested[session] and \
                                len(self.requested[session]) > 1:
                    del self.requested[hasvalue]
                    st = data()
                    st.cmd = 'store'
                    st.kv = {}
                    st.kv[hasvalue] = r[2]
                    for p in self.requested[session]:
                        msgbus.put((sockid, st), 'first')
                break
        else:
            data.result = []
            near = [set(self.findnode(data)) | set(step_result)]
            near.sort(cmp = lambda x,y : cmp(x.peerid ^ data.key, \
                                                    y.peerid ^ data.key))
            near = near[:self.k]
            if self.sessions[session] == near and \
                set(self.requested[session]) == set(near):
                del self.requested[session]
                del self.sessions[session]
                self.result[session] = near
                #try:
                raise dhterror(2)
                #except dhterror, e:
                #	return e.value
            else:
                self.sessions[session] = near
                unrequest = set(near) - set(self.requested[session])
                unrequest.sort(cmp = lambda x, y : cmp(x.peerid ^ \
                                        data.key, y.peerid ^ data.key))
                temp = unrequest[:self.a]
            for p in temp:
                msgbus.put((sockid, data), 'first')
            self.requested[session] += temp
            self.requested[session].sort(cmp = lambda x, y : cmp(x.peerid \
                                        ^ peerid, y.peerid ^ peerid))
            self.requested[session]=self.requested[session][:self.k]

    def reflashkb(self, data):
        '''1 arg, data(cmd, kbid) an int from 0 to 159'''
        random.seed(self.peers)
        data.findid = self.localpeer.peerid ^ random.randint(\
                        pow(2, data.kbid), pow(2, data.kbid + 1) - 1)
        self.nodelookup(data)

    def pong(self, data):
        '''1 arg, (from_peerid, peer session cmd)'''
        if data.from_peerid == self.localpeer.peerid and \
                                        data.session in self.sessions:
            self.sessions[data.session] = time.time()
        else:
            msgbus.put((sockid, data), 'first')
            
    def ping(self, data):
        '''1 arg, data(cmd, from_peerid peer)'''
        if len(self.sessions)  >= 320:
            #try:
            raise dhterror(1)
            #except dhterror, e:
            #	return e.value
        else:
            random.seed(self.peers)
            data.session = random.randint(0, 2**160 - 1)
            data.cmd = 'pong'
            data.from_peerid = self.localpeer.peerid
            self.sessions[data.session] = 0
            start = time.time()
            msgbus.put((sockid, data), 'first')
            time.sleep(5)
            if self.sessions[data.session] != 0:
                respondtime = self.sessions[data.session] - start
                del self.sessions[data.session]
                return respondtime
            else:
                del self.sessions[data.session]
                return 0	# timeout
                            
    def storapair(self, data):
        '''1 arg, data(findid, cmd kv)'''
        try:
            temp = self.nodelookup(data.findid)
        except dhterror, e:
            return e.value
        del data.findid
        data.cmd = 'store'
        for p in temp:
            msgbus.put((sockid, data), 'first')
            
    def joinus(self, data):
        '''1 arg, data.comid .peer .cmd'''
        self.peercome(data)
        data.findid = self.localpeer.peerid
        self.nodelookup(data)
        for i in xrang(160):
            data.kbid = i
            self.reflashkb(data)
            self.buckettimeout[i] = time.time() + self.bucketsecond
            
    def overcaching(self):
        for k,v in self.overcache.items():
            if time.time() >= v:
                del self.dhtcache[k]
                del self.overcache[k]